# 使用 RNN Cell处理语句

import torch

input_size = 4
hidden_size = 4
batch_size = 1

# 词汇表
idx2char = ['e', 'h', 'l', 'o']

# 训练集和样本
x_data = [1, 0, 2, 2, 3]
y_data = [3, 1, 2, 3, 2]

# 独热向量矩阵，用于查找相应的独热向量
one_hot_lookup = [[1, 0, 0, 0],
                  [0, 1, 0, 0],
                  [0, 0, 1, 0],
                  [0, 0, 0, 1]]

# 将输入样本的特征使用独热向量表示
x_one_hot = [one_hot_lookup[x] for x in x_data]
# 修改输入样本和标签的特征维度
inputs = torch.Tensor(x_one_hot).view(-1, batch_size, input_size)
labels = torch.LongTensor(y_data).view(-1, 1)


# 设计模型
class Model(torch.nn.Module):
    def __init__(self, input_size, hidden_size, batch_size):
        super(Model, self).__init__()
        self.batch_size = batch_size
        self.input_size = input_size
        self.hidden_size = hidden_size
        # 定义 RNN Cell 结构
        self.rnncell = torch.nn.RNNCell(input_size=self.input_size,
                                        hidden_size=self.hidden_size)

    # 前馈计算，。传入参数 inputs为输入，hidden为隐层输入
    def forward(self, inputs, hidden):
        hidden = self.rnncell(inputs, hidden)
        return hidden

    # 初始化隐层输入 h0
    def init_hidden(self):
        return torch.zeros(self.batch_size, self.hidden_size)


# 模型实例
model = Model(input_size, hidden_size, batch_size)

# 定义损失和优化器
criterion = torch.nn.CrossEntropyLoss()
# 这里改用一个新的算法 Adam 作为优化算法，学习率改为0.1加速收敛速度
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

# 训练模型
for epoch in range(15):
    loss = 0
    # 梯度清 0
    optimizer.zero_grad()
    # 获得初始化h0
    hidden = model.init_hidden()
    print('Predicted string: ', end='')
    # 这里inputs是要给维度为 (seqLen, batchSize, inputSize) 的输入，所以这里遍历每次去出的元素是序列不同的元素，x1，x2，x3，x4等
    for input, label in zip(inputs, labels):
        # 隐层输出，下一时刻的输入
        hidden = model(input, hidden)
        # 计算损失，这里必须所有元素遍历完，才能作为一轮训练的损失进行反向传播，所以这里需要相加构成计算图
        loss += criterion(hidden, label)
        # 获得对应的具体的值和该值对应的预测的标签
        _, idx = hidden.max(dim=1)
        print(idx2char[idx.item()], end='')
    # 这里循环结束才是一个batch内的所有序列结束，所有序列遍历结束才能进行反向传播和参数更新
    # 反向传播
    loss.backward()
    # 更新参数
    optimizer.step()
    print(', Epoch [%d/15] loss=%.4f' % (epoch + 1, loss.item()))
